﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Raven.Client.Documents;
using Raven.Client.Documents.AI;
using Raven.Client.Documents.Conventions;
using Raven.Client.Documents.Operations.AI;
using Raven.Client.Documents.Operations.AI.Agents;
using Raven.Client.Documents.Operations.Backups;
using Raven.Client.Documents.Operations.ConnectionStrings;
using Raven.Client.Exceptions;
using Raven.Client.Exceptions.Commercial;
using Raven.Client.ServerWide;
using Raven.Client.ServerWide.Operations;
using Raven.Client.Util;
using Raven.Server;
using Raven.Server.Commercial;
using Raven.Server.ServerWide.Commands;
using Sparrow.Json;
using Tests.Infrastructure;
using Xunit;
using Xunit.Abstractions;

namespace SlowTests.Server.Documents.AI.AiAgent;

public class RavenDB_25430 : ReplicationTestBase
{
    public RavenDB_25430(ITestOutputHelper output) : base(output)
    {
    }

    [RavenTheory(RavenTestCategory.Ai, Skip = "RavenDB-25471")]
    [RavenGenAiData(IntegrationType = RavenAiIntegration.OpenAi, DatabaseMode = RavenDatabaseMode.Single, Data = new object[] { BackupType.Backup })]
    [RavenGenAiData(IntegrationType = RavenAiIntegration.OpenAi, DatabaseMode = RavenDatabaseMode.Single, Data = new object[] { BackupType.Snapshot })]
    public async Task DisableAiAgentsOnRestoreWithoutLicense(Options options, GenAiConfiguration aiConfig, BackupType backupType)
    {
        DoNotReuseServer();

        var backupPath = NewDataPath(suffix: "BackupFolder");

        using var source = GetDocumentStore();
        await DisableRevisionCompression(Server, source);

        await source.Maintenance.SendAsync(new PutConnectionStringOperation<AiConnectionString>(aiConfig.Connection));
        var agents = GetAgents(aiConfig);
        foreach (var agentConfig in agents)
        {
            await source.AI.CreateAgentAsync(agentConfig, AiAgentBasics.OutputSchema.Instance);
        }
        var backupOperation = await source.Maintenance.SendAsync(new BackupOperation(new BackupConfiguration
        {
            BackupType = backupType,
            LocalSettings = new LocalSettings
            {
                FolderPath = backupPath
            }
        }));
        await backupOperation.WaitForCompletionAsync(TimeSpan.FromSeconds(30));
        await source.Maintenance.Server.SendAsync(new DeleteDatabasesOperation(source.Database, hardDelete: true));

        await PutLicense(Server, LicenseTestBase.RL_COMM);

        using var store = GetDocumentStore();

        using var destination = new DocumentStore { Urls = new[] { Server.WebUrl }, Database = GetDatabaseName() + "_Restore" }.Initialize();

        foreach (var agentConfig in agents)
        {
            agentConfig.Disabled = true;
        }

        using (Backup.RestoreDatabase(destination,
                   new RestoreBackupConfiguration
                   {
                       BackupLocation = Directory.GetDirectories(backupPath).First(),
                       DatabaseName = destination.Database,
                       DisableOngoingTasks = true
                   }))
        {
            var destRecord = await destination.Maintenance.Server.SendAsync(new GetDatabaseRecordOperation(destination.Database));
            var destConfigs = destRecord.AiAgents;
            Assert.NotNull(destConfigs);
            Assert.Equal(2, destConfigs.Count);
            Assert.NotNull(destConfigs[0]);
            Assert.NotNull(destConfigs[1]);
            Assert.Equal("shopping-assistant", destConfigs[0].Identifier);
            Assert.Equal("warehouse-manager", destConfigs[1].Identifier);
            using (var context = JsonOperationContext.ShortTermSingleUse())
            {
                var converter = DocumentConventions.Default.Serialization.DefaultConverter;
                var c0 = converter.ToBlittable(agents[0], context);
                var d0 = converter.ToBlittable(destConfigs[0], context);
                Assert.Equal(c0, d0);
                var c1 = converter.ToBlittable(agents[1], context);
                var d1 = converter.ToBlittable(destConfigs[1], context);
                Assert.Equal(c1, d1);
            }
        }
    }

    [RavenTheory(RavenTestCategory.Ai, Skip = "RavenDB-25471")]
    [RavenGenAiData(IntegrationType = RavenAiIntegration.OpenAi, DatabaseMode = RavenDatabaseMode.Single, Data = new object[] { BackupType.Backup })]
    [RavenGenAiData(IntegrationType = RavenAiIntegration.OpenAi, DatabaseMode = RavenDatabaseMode.Single, Data = new object[] { BackupType.Snapshot })]
    public async Task ThrowAiAgentsOnRestoreWithoutLicense(Options options, GenAiConfiguration aiConfig, BackupType backupType)
    {
        DoNotReuseServer();

        var backupPath = NewDataPath(suffix: "BackupFolder");

        using var source = GetDocumentStore();
        await DisableRevisionCompression(Server, source);

        await source.Maintenance.SendAsync(new PutConnectionStringOperation<AiConnectionString>(aiConfig.Connection));
        var agents = GetAgents(aiConfig);
        foreach (var agentConfig in agents)
        {
            await source.AI.CreateAgentAsync(agentConfig, AiAgentBasics.OutputSchema.Instance);
        }
        var backupOperation = await source.Maintenance.SendAsync(new BackupOperation(new BackupConfiguration
        {
            BackupType = backupType,
            LocalSettings = new LocalSettings
            {
                FolderPath = backupPath
            }
        }));
        await backupOperation.WaitForCompletionAsync(TimeSpan.FromSeconds(30));
        foreach (var agentConfig in agents)
        {
            agentConfig.Disabled = true;
            await source.AI.CreateAgentAsync(agentConfig, AiAgentBasics.OutputSchema.Instance);
        }

        await PutLicense(Server, LicenseTestBase.RL_COMM);

        using var destination = new DocumentStore { Urls = new[] { Server.WebUrl }, Database = GetDatabaseName() + "_Restore" }.Initialize();
        var e = Assert.Throws<LicenseLimitException>(() => {
            using (Backup.RestoreDatabase(destination,
                       new RestoreBackupConfiguration
                       {
                           BackupLocation = Directory.GetDirectories(backupPath).First(),
                           DatabaseName = destination.Database,
                       }))
            {
            }
        });
        Assert.Contains("Your license doesn't support using the AI Agent feature", e.Message);

        e = await Assert.ThrowsAsync<LicenseLimitException>(async () => {
            foreach (var agentConfig in agents)
            {
                agentConfig.Disabled = false;
                await source.AI.CreateAgentAsync(agentConfig, AiAgentBasics.OutputSchema.Instance);
            }
        });
        Assert.Contains("Your current license doesn't include the AI Agent feature", e.Message);

        e = await Assert.ThrowsAsync<LicenseLimitException>(async () => {
            var record = source.Maintenance.Server.Send(new GetDatabaseRecordOperation(source.Database));
            record.AiAgents[0].Disabled = false;
            record.DatabaseName += "_Copy";
            await source.Maintenance.Server.SendAsync(new CreateDatabaseOperation(record));
        });
        Assert.Contains("Your license doesn't support using the AI Agent feature", e.Message);

        var chat = source.AI.Conversation(agents[0].Identifier, "chats/",
            new AiConversationCreationOptions()
                .AddParameter("company", "companies/90-A"));

        var ex = await Assert.ThrowsAsync<RavenException>(async () => {
            chat.SetUserPrompt("hello");
            await chat.RunAsync<AiAgentBasics.OutputSchema>(CancellationToken.None);
        });
        Assert.Contains("The AI Agent 'shopping-assistant' is currently disabled. Please enable the agent before starting\\continuing a conversation.", ex.Message);
    }

    private static List<AiAgentConfiguration> GetAgents(GenAiConfiguration aiConfig)
    {
        var agent0 = new AiAgentConfiguration("shopping assistant", aiConfig.ConnectionStringName,
            "You are an AI agent of an online shop, helping customers answer queries about that topic only. When talking about orders or products, include the ids as well.");
        agent0.Identifier = "shopping-assistant";
        agent0.Parameters.Add(new AiAgentParameter("company"));
        agent0.ChatTrimming = null;
        agent0.Queries =
        [
            new AiAgentToolQuery
            {
                Name = "ProductSearch",
                Description = "semantic search the store product catalog",
                Query = "from Products where vector.search(embedding.text(Name), $query)",
                ParametersSampleObject = "{\"query\": [\"term or phrase to search in the catalog\"]}"
            },
            new AiAgentToolQuery
            {
                Name = "RecentOrder",
                Description = "Get the recent orders of the current user",
                Query = "from Orders where Company = $company order by OrderedAt desc limit 10",
                ParametersSampleObject = "{}"
            }
        ];


        var agent1 = new AiAgentConfiguration("warehouse manager", aiConfig.ConnectionStringName, "You are an AI agent managing a warehouse.");
        agent1.Identifier = "warehouse-manager";
        agent1.Actions =
        [
            new AiAgentToolAction
            {
                Name = "ProductSearch",
                Description = "semantic search the store product catalog",
                ParametersSampleObject = "{\"query\": [\"term or phrase to search in the catalog\"]}"
            },
            new AiAgentToolAction
            {
                Name = "RecentOrder",
                Description = "Get the recent orders of the current user",
                ParametersSampleObject = "{}"
            }
        ];
        agent1.ChatTrimming = null;
        return new List<AiAgentConfiguration>() { agent0, agent1 };
    }

    private static async Task PutLicense(RavenServer leader, string licenseType)
    {
        var license = Environment.GetEnvironmentVariable(licenseType);
        Raven.Server.Commercial.LicenseHelper.TryDeserializeLicense(license, out License li);

        await leader.ServerStore.PutLicenseAsync(li, RaftIdGenerator.NewId());
    }

    private static async Task DisableRevisionCompression(RavenServer leader, DocumentStore store)
    {
        var command = new EditDocumentsCompressionCommand(new DocumentsCompressionConfiguration { CompressRevisions = false, Collections = new string[] { } }, store.Database,
            RaftIdGenerator.NewId());
        await leader.ServerStore.SendToLeaderAsync(command);
    }
}
